import glob
import random

import numpy as np
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import os
from PIL import Image
import json
from pycocotools import mask as mask_utils
import torch
from torch.nn import functional as F
from tqdm import tqdm

class MultiTaskDataset(Dataset):
    def __init__(self, root: str, split: str = "train", transform=None, image_size=256,
                 v_patch_nums=(1, 2, 3, 4, 5, 6, 8, 10, 13, 16), separator=False, **kwargs):
        self.transforms = transform
        self.split = split
        self.image_size = image_size
        self.v_patch_nums = v_patch_nums
        self.separator = separator
        
        # 只加载GoPro数据集路径
        self.gopro_root = os.path.join(root)
        self.blur_paths = sorted(glob.glob(os.path.join(self.gopro_root, split, "input", "*.png")))
        self.sharp_paths = sorted(glob.glob(os.path.join(self.gopro_root, split, "target", "*.png")))
        
        print(f'Found {len(self.blur_paths)} images in GoPro {split} set')
        
        self.default_cls = 0
        
        # 任务ID映射 - 可以根据需要添加更多任务
        self.task_ids = {
            'deblur': 0,          # 去模糊任务
            'color': 1,           # 灰度图上色
            'deblur_color': 2,   # 灰度图去模糊
        }

    def __len__(self):
        return len(self.blur_paths)

    def __getitem__(self, index: int):
        # 随机选择任务类型
        task_type = random.choices(['deblur', 'color', 'deblur_color'], [0.33, 0.33, 0.33], k=1)[0]
        
        # 加载模糊和清晰图像
        blur_path = self.blur_paths[index]
        sharp_path = self.sharp_paths[index]
        
        blur_img = Image.open(blur_path).convert('RGB')
        sharp_img = Image.open(sharp_path).convert('RGB')
        
        if task_type == 'deblur':
            # 普通去模糊任务
            if self.transforms:
                sharp_img, blur_img = self.transforms(sharp_img, blur_img)
            
            image = sharp_img
            condition = blur_img
            
        elif task_type == 'color':  # gray color task
            # 灰度图上色任务
            # 将清晰图像转换为灰度图
            sharp_img_gray = sharp_img.convert('L').convert('RGB')
            
            if self.transforms:
                sharp_img, sharp_img_gray = self.transforms(sharp_img, sharp_img_gray)
            
            image = sharp_img
            condition = sharp_img_gray
        elif task_type == 'deblur_color':
            # 复合任务：将模糊图转为灰度
            blur_gray_img = blur_img.convert('L').convert('RGB')
            
            if self.transforms:
                sharp_img, blur_gray_img = self.transforms(sharp_img, blur_gray_img)
            
            image = sharp_img
            condition = blur_gray_img
        
        
        ignore_masks = torch.ones((1378,)) if self.separator else torch.ones((1360,))
        ignore_masks_ = torch.ones((1378,)) if self.separator else torch.ones((1360,))
        
        return {
            'image': image,
            'mask': condition,
            'cls': self.default_cls,
            'type': torch.tensor(self.task_ids[task_type]),
            'ignore_mask': ignore_masks,
            'ignore_mask_': ignore_masks_
        }


 